#!/usr/bin/env python

import ConfigParser,subprocess,os,sys,time,glob,shutil,select
import multiprocessing,signal,random,copy,hashlib,datetime,gzip

config = ConfigParser.ConfigParser()

try:
	dirname  = os.path.dirname(os.path.abspath(sys.argv[0]))	
	config.read(dirname + '/settings.cfg')
except Exception,err:
	print err
	print "Unable to read settings.cfg"
	sys.exit(-1)
	raise

total_nodes = int(config.get('general','total_nodes'))
duration = int(config.get('general','experiment_duration'))
initial_wait = int(config.get('general','initial_wait'))
shutdown_wait = int(config.get('general','shutdown_wait'))

random.seed(int(config.get('general','initial_random_seed')))

show_debug = config.getboolean('general','show_debug_messages')

use_chroot = config.getboolean('general','use_chroot')

log_name = 'core'

green = '\033[01;32m'
red = '\033[01;31m'
defaultColor = '\033[0;0m'

environment = (total_nodes+1) * [None]
for n in range(0,total_nodes+1):
	environment[n] = {}

for k in os.environ:
	environment[0][k] = os.environ[k]

for k in config.items('general'):
	environment[0][k[0]] = k[1]

nodeid = "0"

logfile = None

def open_logfile(filename):
	global logfile
	if not logfile == None:
		logfile.close()
	logfile = open(filename,'w')

def close_logfile():
	global logfile
	if not logfile == None:
		logfile.close()
		logfile = None

# LOGGING #
def logger(tag,message):
	global log_name
	color = None
	if tag == "info":
		color = green
	elif tag == "error":
		color = red

	s = '%s:%s:%s:%s' % (log_name,time.time(),tag,str(message).strip())

	if color:
		print color + s + defaultColor
	else:
		print s

	if logfile:
		logfile.write(s + '\n')
		logfile.flush()

def log(s):
	logger('info',s)

def log_debug(s):
	if show_debug:
		logger('debug',s)

def log_err(s):
	logger('error',s)

def log_cmd(command,s):
	logger('cmd',command + ':' + s)

	if (command == "./waf" and s.split(':')[0] == "drift"):
		drift = s.split(':')
		real_time = int(drift[1].strip())
		sim_time = int(drift[2].strip())
		delta = int(drift[3].strip())
		if delta > 100:
			log_err("Drift is more than 100 ms (%s). May trigger congestion event." % delta)

def log_cmd_err(command,s):
	logger('cmd_error', command + ":" + s)

def start_nodes(experiment):
	global total_nodes,config,use_chroot
	
	save_env()

	for n in range(1,total_nodes+1):
		log("starting lxc node %s" % (n))
#		safe_exec_fork(['lxc-start','--name','node'+str(n),'./run.py','startnode',str(n),experiment])
		if use_chroot:
			script = config.get('directories','chroot_bindpoint') + "/run.py"
		else:
			script = "./run.py"
		subprocess.Popen(['lxc-execute','--name','node'+str(n),script,'startnode',str(n),experiment],close_fds=True)

def store_dict(fname, d):
	f = file(fname,'w')
	f.write(repr(d))
	f.close()

def read_dict(fname):
	f = file(fname,'r')
	return eval(f.read())

def save_env():
	global config,nodeid,use_chroot
	if (int(nodeid) == 0): # node id 0 is core
		store_dict(config.get('directories','dumps')+"/lxc-environment.dat", environment)
	else:
		if use_chroot:
			store_dict(config.get('directories','chroot_working_dir')+"/lxc-environment.dat", environment)
		else:
			raise Exception("Warning: lxc container attempted to overwrite global lxc-environment.dat")

def read_env():
	global environment,nodeid
	if (int(nodeid) == 0): # node id 0 is core
		environment = read_dict(config.get('directories','dumps')+"/lxc-environment.dat")
	else:
		if use_chroot:
			environment = read_dict(config.get('directories','chroot_working_dir')+"/lxc-environment.dat")
		else:
			environment = read_dict(config.get('directories','dumps')+"/lxc-environment.dat")	

def chroot_mount():
	global config
	mobiemu_in_chroot = config.get('directories','chroot_bindpoint')
	for rootpoint in config.get('directories','chroot_rootpoints').split(','):
		for mountpoint in glob.glob(rootpoint):
			log("mounting with --bind in %s" % mountpoint)
			safe_exec(['mount','--bind',config.get('directories','main'),mountpoint+"/"+mobiemu_in_chroot],chdir=False)

def chroot_umount():	
	global config
	mobiemu_in_chroot = config.get('directories','chroot_bindpoint')
	for rootpoint in config.get('directories','chroot_rootpoints').split(','):
		for mountpoint in glob.glob(rootpoint):
			log("umounting %s" % mountpoint)
			safe_exec_ignore_errors(['umount',mountpoint + "/" + mobiemu_in_chroot],chdir=False)

# replaces @NN@ with node id, @HH@ with node id in hex, @MAC@ with mac address and @MTU@ with MTU
def replace_string(s, nodeid, mac, mtu):	
	ret = s.replace('@NN@',str(nodeid))
	ret = ret.replace('@HH@',hex(nodeid))
	ret = ret.replace('@MAC@',mac)
	ret = ret.replace('@MTU@',mtu)

	if multiprocessing.cpu_count() > 1: # By default, schedule on all but last core
		cpu_core = "0-%s" % (multiprocessing.cpu_count()-2)
	else:
		cpu_core = "0"

	ret = ret.replace('@CPUS@',cpu_core)
	return ret


def create_lxc_config(lxcconfig,nodeid,mac,mtu):
	global config
	output_name = lxcconfig
	newfile = open(output_name,'w')
	template = open(config.get('directories','configs')+'/lxc.conf','r')

	for s in template:
		if len(s.lstrip()) > 0 and s.lstrip()[0] == '#':
			continue
		s = replace_string(s, nodeid, mac, mtu)
		newfile.write(s)

	newfile.close
	template.close

	return output_name

# execute and throw an exception on error 
def safe_exec(cmd,simulator_core=False,use_shell=False,chdir=False):
	global config,nodeid,environment

	commands = copy.copy(cmd) 

	log_prefix = commands[0]

	if chdir:
		if os.path.exists(commands[0]):
			needs_prefix = True
		else:
			needs_prefix = False

		if use_chroot:
			dumps_dir = config.get('directories','chroot_working_dir')
		else:
			dumps_dir = config.get('directories','dumps')

		old_dir = os.getcwd()
		os.chdir(dumps_dir)

		if needs_prefix:
			commands[0] = old_dir + '/' + commands[0]

		log_debug("executing command in %s" % dumps_dir)

	# use last core for simulator, the others for normal tasks
	if simulator_core == False: 
		if multiprocessing.cpu_count() > 1: # By default, schedule on all but last core
			cpu_core = "0-%s" % (multiprocessing.cpu_count()-2)
		else:
			cpu_core = "0"
		niceness = config.get('general','experiment_niceness')
		log_debug("Executing on simulator core")
	else:
		# schedule simulator tasks on last core
		cpu_core = "%s" % (multiprocessing.cpu_count()-1) 
		niceness = config.get('general','simulator_niceness')

	commands[0:0] = ['nice','-n',niceness]
	commands[0:0] = ['taskset','-c',cpu_core]

	if (nodeid != 0):
		_env = environment[int(nodeid)]
	else:
		_env = {}
	_env.update(environment[0])

	# run shell commands as strings, not dicts
	if (use_shell):
		commands = ' '.join(commands)

	log_debug("%s" % commands)

	try:
		child = subprocess.Popen(commands,stdout=subprocess.PIPE, bufsize=1, stderr=subprocess.PIPE, close_fds=True, shell=use_shell,env=_env)
	
	finally:
		if chdir:
			os.chdir(old_dir)

	rs = [child.stdout,child.stderr]

	while rs:
		try:
			rlist,wlist,xlist = select.select(rs,[],[])
		except select.error, e:
			if e.args[0] == error.EINTR:
				continue
			raise

		if child.stdout in rlist:
			line = child.stdout.readline()
			if line == "":
				child.stdout.close()
				rs.remove(child.stdout)
			else:
				log_cmd(log_prefix,line)

		if child.stderr in rlist:
			line = child.stderr.readline()
			if line == "":
				child.stderr.close()
				rs.remove(child.stderr)
			else:
				log_cmd_err(log_prefix,line)

	child.wait()		

	if child.poll() == None:
		raise Exception("Command still running after wait(). This should not happen.")

	if child.poll() != 0:
		raise Exception("Command exited with error %s" % child.poll())

def safe_exec_ignore_errors(commands,simulator_core=False,chdir=False):
	try:
		safe_exec(commands,simulator_core=simulator_core,chdir=chdir)
	except Exception,err:
		log_debug("%s (ignored)" % err)
		pass

def safe_exec_fork(commands,simulator_core=False,chdir=False,use_shell=False): 
	process = multiprocessing.Process(target=safe_exec, args=(commands,simulator_core,use_shell,chdir,))
	process.start()
	return process

def create_devices(mtu):
	global total_nodes,config,use_chroot
	
	try:
		if use_chroot:
			chroot_mount()
	except:
		log_err("got error when mounting in chroot-containers")
		raise

	if not os.path.ismount("/sys/fs/cgroup"):
		log("mounting /cgroup")
		if not os.path.isdir('/cgroup'):
			os.mkdir('/cgroup')
	
		if not os.path.ismount('/cgroup'):
			safe_exec(['mount','-t','cgroup','cgroup','/cgroup'])
	
	log("creating devices with mtu %s" % mtu)

	for n in range(1,total_nodes+1):
		tapdev='tap-'+str(n)
		brdev='br-'+str(n)
		node='node'+str(n)
		mac='00:16:3e:00:01:%0.2X' % int(n)

		log( "creating bridge %s and tap-device %s for node %s with mac %s" % (brdev,tapdev,node,mac))

		lxcconfig='/tmp/lxc-'+str(n)+'.conf'

		# set up virtual devices
		safe_exec(['brctl','addbr',brdev])
		safe_exec(['tunctl','-t',tapdev])
		safe_exec(['ifconfig',tapdev,'0.0.0.0','mtu',str(mtu),'promisc','up'])
		safe_exec(['brctl','addif',brdev,tapdev])
		safe_exec(['ifconfig',brdev,'up'])
		safe_exec(['ifconfig',brdev,'mtu',mtu,'hw','ether',mac])

		# write env
		environment[n]['tap_device'] = tapdev
		environment[n]['bridge_device'] = brdev
		environment[n]['node'] = node
		environment[n]['node_id'] = str(n)
		environment[n]['mac_address'] = mac

		# create lxc node
		create_lxc_config(lxcconfig,n,mac,mtu)
		safe_exec(['lxc-create','-n',node,'-f',lxcconfig])

	for bnf in glob.glob('/proc/sys/net/bridge/bridge-nf-*'):
		log("echo 0 > %s" % bnf)
		f = open(bnf,'w')
		f.write('0')
		f.close

def remove_devices():
	global total_nodes,use_chroot

	try:
		if use_chroot:
			chroot_umount()
	except:
		log_err("got error when umounting in chroot-containers")
		raise

	for n in range(1,total_nodes+1):
		tapdev='tap-'+str(n)
		brdev='br-'+str(n)
		node='node'+str(n)
		lxcconfig='/tmp/lxc-'+str(n)+'.conf'

		log("removing bridge %s and tap device %s for node %s" % (brdev,tapdev,node))

		# stop lxc
		safe_exec_ignore_errors(['lxc-stop','-n',node])
		safe_exec_ignore_errors(['lxc-kill','-n',node])
		safe_exec_ignore_errors(['lxc-destroy','-n',node])

		# remove bridge and tap device
		safe_exec_ignore_errors(['ifconfig',brdev,'down'])
		safe_exec_ignore_errors(['brctl','delif',brdev,tapdev])
		safe_exec_ignore_errors(['brctl','delbr',brdev])
		safe_exec_ignore_errors(['ifconfig',tapdev,'down'])
		safe_exec_ignore_errors(['tunctl','-d',tapdev])

		# unlink config file
		try:
			if os.path.exists(lxcconfig):
				os.unlink(lxcconfig)
		except:
			log_err("Unable to unlink %s" % lxcconfig)

	# remove any remaining links
	links = subprocess.Popen(['ip','link','show'],stdout=subprocess.PIPE,close_fds=True)
	for link in links.stdout:
		if link.find('veth') > 0 or link.find('tap-') > 0:
			device = link.split(':')[1].strip()
			log("Device  %s still up. Removing..." % device)
			safe_exec_ignore_errors(['ip','link','del',device]) # no point in using safe_exec here

	if os.path.ismount('/cgroup'):
		log("umounting /cgroup")
		safe_exec_ignore_errors(['umount','/cgroup'])
		if os.path.exists('/cgroup'):
			try:
				os.rmdir('/cgroup')
			except:
				log_err("Unable to remove /cgroup")



# should be run outside lxc
def start_ns3(parameters):
	global config, ns3_pid, total_nodes, duration, initial_wait, shutdown_wait, rand_seed
	ns3_pid = None

	log("Starting ns3")

	ns3_path = os.path.abspath(config.get('ns3','path'))
	if not os.path.isdir(ns3_path) or not os.path.exists(ns3_path + '/waf'):
		raise Exception("ns3 not found in " + ns3_path + ". Change ns3.path in settings.cfg to point to your ns-3 installation.")

	scratch_path = ns3_path + '/scratch/'

	if not os.path.isdir(scratch_path):
		raise Exception("scratch directory not found in "+ns3_path)

	script = os.path.abspath(config.get('directories','topologies')) + '/' + parameters['script']

	if not os.path.exists(script):
		raise Exception("simulation script not found in " + script)

	shutil.copy(script,scratch_path)

	cwd = os.getcwd()
	os.chdir(ns3_path)
	try:
		safe_exec(['./waf','-d','optimized'],chdir=False) # compile first
	except:
		log_err("ns3 compilation error. Exiting")
		raise

	# TODO remove extension properly, not just :-3
	command = ['./waf','-d','optimized','--run']
	s = "scratch/" + parameters['script'][:-3]
	for c in parameters.items():
		if c[0] != "script":		 
			s = s + ' --'+ c[0] + "=" + c[1]

	s = s + ' --nodes=%s' % total_nodes
	s = s + ' --seed=%s' % rand_seed 
	s = s + ' --duration=%s' % (duration+initial_wait+shutdown_wait)

	command.append(s)
	log("ns3 command: %s" % command)


	ns3_pid = safe_exec_fork(command, True,chdir=False)
	os.chdir(cwd)
	#... ns3_pid = subprocess.Popen(

def stop_ns3():
	global ns3_pid
	if ns3_pid != None:
		log("Stopping ns3...")
		ns3_pid.terminate()
		if ns3_pid.is_alive:
			time.sleep(1)
		if ns3_pid.is_alive:
			os.kill(ns3_pid.pid,signal.SIGTERM)

# should be run from within lxc
def start_module(nodeid, module):
	global config
	path=config.get('directories','modules') + '/' + module
	log("starting module %s" % path)

	log_debug("setting environment variables for nodeid %s" % nodeid)
	try:
		for k in config.items(module):
	#		safe_exec(['/bin/sh export',k[0]+'="' + k[1] + '"'],use_shell=True)
			environment[int(nodeid)][k[0]] = str(k[1])
	except Exception,err:
			log_err(err)

	log_debug("executing module")
	try:
		safe_exec_fork([path,'start'],chdir=True)
	except:
		log_err("unable to start module %s" % path)
		raise

	# remember correct core

# should be run within lxc
def start_all_modules(nodeid, modules):
	for mod in modules:
		start_module(nodeid, mod)

def stop_all_modules(modules):
	path = config.get('directories','modules')
	for mod in modules:
		try:
			safe_exec([path + '/' + mod, 'stop'],chdir=True)
		except:
			log_err("unable to stop module %s in path %s" % (mod,path))
			raise

# this should be run from within lxc
def start_experiment(experiment):
	global config,environment
	log("starting experiment %s" % experiment)

	log_debug("setting environment variables for experiment %s" % experiment)
	try:
		for k in config.items(experiment):
			environment[0][k[0]] = str(k[1])
	except Exception,err:
			log_err(err)

	start_nodes(experiment)	

# create archive
def create_archive():
	archive = config.get('directories','dumps')+'/mobiemu_config_archive.tar.gz'
	include = config.get('source_archiver','include')

	log("archiving configuration files in %s..." % archive)

	command ="tar --preserve-order cvzf " + archive + " `find " + include + " | sort --stable`"
	safe_exec(['tar','cvzf',archive,include],use_shell=True,chdir=False)

	return archive

def md5_file(filename):
	f = gzip.open(filename)
	md5 = hashlib.md5()
	while True:
		data = f.read(8192)
		if not data:
			break
		md5.update(data)
	f.close()
	return md5.hexdigest()

def get_date_string():
	ds = str(time.ctime())
	ds = ds.replace(' ','_')
	ds = ds.replace(':','')
	ds = ds.replace('/','')
	return ds

def move_results(directory_tag=None):
	log("moving results to archive directory")
	archive = create_archive()
	log("calculating source checksum...")
	checksum = md5_file(archive)
	log(checksum)
	archive_md5 = config.get('general','experiments').replace(',','_') + "_" + config.get('general','topology') + "_" + checksum
	

	# move results in to the following directory structure
	# results/md5 of source archive/directory_tag

	results_base_dir = config.get('directories','results')
	if not os.path.exists(results_base_dir):
		raise Exception("Results directory (%s) doesn't exist." % results_base_dir)

	results_dir = results_base_dir + "/" + archive_md5
	if not os.path.exists(results_dir):
		log("results directory for this source checksum not found, creating...")
		os.mkdir(results_dir)
		shutil.move(archive,results_dir)

	if directory_tag == None:
		directory_tag = get_date_string()
	else:
		directory_tag = directory_tag + "_" + get_date_string()

	target_dir = results_dir + "/" + directory_tag
	
	if os.path.exists(target_dir):
		raise Exception("target directory already exists (%s). Aborting..." % target_dir)

	os.mkdir(target_dir)

	include = config.get('results_archiver','include')
	try:
		safe_exec(['mv','-v',include,target_dir],use_shell=True,chdir=False)
	except:
		log_err("unable to move results to " + target_dir)
		raise

	log("results moved from %s to %s" % (include,target_dir))
		
# Split all parameters which have a list of options separated by ,
# For each option, a new parameter set is created and finally
# a set with all possible combinations is returned.
def permute_var(items):
	ret = []
	permuted = False
	for i in range(len(items)):
		if items[i][1].find(',') > -1:
			permuted = True
			for val in items[i][1].split(','):
				new_items = items[:]
				new_items[i] = (new_items[i][0],val)
				for stuff in permute_var(new_items):
					ret.append(stuff)
	if not permuted:	
		ret.append(items)
	return ret
		

total_exp = 0
fake_run = True

# execute within lxc node
if len(sys.argv) > 1:
	if sys.argv[1] == 'startnode':
		nodeid = sys.argv[2]
		log_name = "node" + nodeid
		experiment_path = config.get('directories','experiments')
		experiment_script = config.get(sys.argv[3],'experiment')
		modules = config.get(sys.argv[3],'modules').split(',')

		log("Starting node %s" % nodeid)

		orig_dir = os.getcwd()

		if use_chroot:
			chroot_bindpoint = config.get('directories','chroot_bindpoint')
			log("in chroot, changing directory to %s" % chroot_bindpoint)
			orig_dir = os.getcwd()
			os.chdir(chroot_bindpoint)

		open_logfile(config.get('directories','dumps') + "/" + log_name + ".log")

		read_env()

		try:
			log("Starting modules...")
			start_all_modules(nodeid, modules)

			log("Waiting for %s" % initial_wait)
			time.sleep(initial_wait)

			log("Running experiment %s" % experiment_script)
			process = safe_exec_fork([experiment_path+'/'+experiment_script,nodeid],chdir=True)
			time.sleep(duration)

			log("Waiting for shutdown...")
			time.sleep(shutdown_wait/2)

			if process.is_alive():
				log_err("Process has not terminated after %s seconds. Sending SIGTERM" % (shutdown_wait/2))
				process.terminate()
				time.sleep(shutdown_wait/4)

			if process.is_alive():
				log_err("Process still running. Attempting SIGKILL.")
				os.kill(process.pid,signal.SIGKILL)
				time.sleep(shutdown_wait/4)

			if process.is_alive():
				log_err("Unable to terminate experiment. Exiting anyway.")
			else:
				log("Experiment is no longer running.")

			log("Stopping modules...")
			stop_all_modules(modules)

			log("done, exiting")
		except Exception,err:
			log_err(str(err))
			log_err("NODE CRASHED - STOPPING MODULES")
			stop_all_modules(modules)
			raise

		if use_chroot:
			os.chdir(orig_dir)

		close_logfile()

		sys.exit()

	# cleanup
	if sys.argv[1] == 'cleanup':
		log("Cleaning up...")
		remove_devices()		
		sys.exit()

	# archive
	if sys.argv[1] ==  'archive':
		create_archive()
		sys.exit()

	# moveresults
	if sys.argv[1] == 'moveresults':
		move_results()
		sys.exit()

	# run
	if sys.argv[1] == 'run':
		log("Running experiment...")
		fake_run = False

if not fake_run:
	if use_chroot:
		log("Using chroot.")
	
	open_logfile(config.get('directories','dumps') + "/core_startup.log")

	# Archive configuration files
	create_archive()

	log("running init scripts")
	for script in config.get('general','init_scripts').split(','):
		log(script)
		safe_exec([config.get('directories','modules')+"/"+script])
	
# Create all parameter combinations for the experiments
for rep in range(int(config.get('general','repetitions'))):
	rand_seed = random.randint(0,65535) # set new rand seed every repetition
	for experiment in config.get('general','experiments').split(','):
		for topology in config.get('general','topology').split(','):
			for f in permute_var(config.items(topology)):
				if not fake_run:

					open_logfile(config.get('directories','dumps') + "/core.log")

					# convert f to dict
					sim_params = {}
					for t in f:
						sim_params[t[0]] = t[1]

					log("%s,%s,%s,%s" % (rep,experiment,topology,sim_params))
					log("---Experiment parameters---")
					environment[0]['current_experiment'] = experiment
					log("experiment=%s" % experiment)
					environment[0]['current_repetition'] = str(rep)
					log("current_repetition=%s" % rep)
					environment[0]['current_topology'] = topology
					log("current_topology=%s" % topology)
					#environment[0]['current_rand_seed'] = str(rand_seed)
					log("current_rand_seed=%s" % rand_seed)

					state = random.getstate()
					random.seed(rand_seed)
					for n in range(1,total_nodes+1):
						environment[n]['current_rand_seed'] = str(random.randint(0,65535))
						log("node "+str(n)+" rand_seed=" + str(environment[n]['current_rand_seed']))
					random.setstate(state)

					log("---Topology parameters---")


					for kv in sim_params.items():
						log("%s=%s" % (kv[0], kv[1]))
						environment[0][kv[0]] = kv[1]

					try:
						create_devices(sim_params['mtu'])
					except:
						log_err('Unable to create virtual devices and linux containers. Try cleanup.')
						raise

					try:
						start_ns3(sim_params)
						start_experiment(experiment)
						time.sleep(initial_wait + duration + shutdown_wait)
					finally:
						stop_ns3()
						remove_devices()
					
					try:
						move_results(experiment + "_" + str(rand_seed))
					except:
						log_err("unable to store results! aborting...")
						raise
	

				total_exp += 1 # calculate number of experiments

	close_logfile() # close any open log files

if fake_run:
	estimated_total_time = total_exp * (duration + initial_wait + shutdown_wait)
	print "Total number of experiments:", total_exp
	print "This would approximately take", estimated_total_time,  "seconds, or", estimated_total_time/3600.0, "hours"
	print
	print "Use ./run.py run to start the experiment."
	print
	print "Currently, run.py must be executed with root privileges."
	print 
	print "Other commands:"
	print "	./run.py cleanup	Remove and stop all lxc-containers, bridges and tap devices"
	print "	./run.py archive	Archive configuration files and store them in ",config.get('directories','dumps')
	print "	./run.py moveresults	Move results from ",config.get('directories','dumps')
